import komand
from .schema import GetMalwareTestFileInput, GetMalwareTestFileOutput
# Custom imports below
import base64


class GetMalwareTestFile(komand.Action):
    def __init__(self):
        super(self.__class__, self).__init__(
                name='get_malware_test_file',
                description='Gets a unique, benign malware test file that will trigger',
                input=GetMalwareTestFileInput(),
                output=GetMalwareTestFileOutput())

    def run(self, params={}):
        client = self.connection.client
        out = base64.b64encode(client.get_malware_test_file()).decode()
        return { 'file': out }

    def test(self):
        """TODO: Test action"""
        client = self.connection.client
        out = base64.b64encode(client.get_malware_test_file()).decode()
        return { 'file': out }
